"""基础神经网络词嵌入投影模型"""
import pickle
from config import *
from tqdm import tqdm
import numpy as np
import torch
import torch.nn as nn
from DataSet import *
from torch.utils.data import DataLoader

torch.manual_seed(1)
with open(Config.word2vec_path, 'rb') as f:
    Word2Vec = pickle.load(f)

class NPM(nn.Module):
    """
    基础神经网络词嵌入投影模型类，
    是一个简单的前馈神经网络，
    单隐藏层，隐藏层维度为100
    """
    def __init__(self):
        super(NPM,self).__init__()
        self.embeds = nn.Embedding.from_pretrained(torch.from_numpy(Word2Vec.word2vec), freeze=True)
        self.fc1 = nn.Linear(100,100)
        self.fc2 = nn.Linear(100,200)

    def forward(self, x):
        x = self.embeds(x).float()
        x = torch.tanh(self.fc1(x))
        out = self.fc2(x)
        return out

